"""
根据配准修改遗传算法的思路
1 种群初始化：已知对X光的估计值和真值之间的误差之间。可以以估计值为中心，不同维度的误差范围内生成初始种群
2 交叉：有理数交叉
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

DNA_SIZE = 24
POP_SIZE = 80
CROSSOVER_RATE = 0.6
MUTATION_RATE = 0.01
N_GENERATIONS = 100
X_BOUND = [-2.048, 2.048]
Y_BOUND = [-2.048, 2.048]


def F(x, y):
    return 100.0 * (y - x ** 2.0) ** 2.0 + (1 - x) ** 2.0  # 以香蕉函数为例


def plot_3d(ax):
    X = np.linspace(*X_BOUND, 100)
    Y = np.linspace(*Y_BOUND, 100)
    X, Y = np.meshgrid(X, Y)
    Z = F(X, Y)
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    plt.pause(3)
    plt.show()


def get_fitness(pop):
    x, y = translateDNA(pop)
    pred = F(x, y)
    return pred
    # return pred - np.min(pred)+1e-3  # 求最大值时的适应度
    # return np.max(pred) - pred + 1e-3  # 求最小值时的适应度，通过这一步fitness的范围为[0, np.max(pred)-np.min(pred)]


def crossover(pop, cross_rate=0.8):  # pop表示种群矩阵，一行表示一个二进制编码表示的DNA，矩阵的行数为种群数目
    # 第1到(k-1)位不变，第k位进行SBX计算，从第k+1位开始做交换
    new_pop = []
    for father in pop:  # 遍历种群中的每一个个体，将该个体作为父亲
        child = father
        if np.random.rand() < cross_rate:  # 产生子代时不是必然发生交叉，而是以一定的概率发生交叉
            mother = pop[np.random.randint(0, population_size)]  # 再种群中选择另一个个体，并将该个体作为母亲
            k = np.random.randint(0, population_size)
            x_pop = child[0:k]  # 将子代的0到k-1位表示X
            y_pop = child[k+1:]  # 将子代的k+1到最后一位表示Y
            # 使用SBX计算第k位
            nc = 20
            u = np.random.rand()
            if u <= 0.5:
                beta = pow(2 * u, 1 / (nc + 1))
            else:
                beta = pow(1 / (2 * (1 - u)), 1 / (nc + 1))
            _child_k = 0.5 * (1 + beta) * _father2[:2] + (1 - beta) * _father2[2:]
            _child2 = np.hstack((_child21, _child22))

    x = x_pop.dot(2 ** np.arange(DNA_SIZE)[::-1]) / float(2 ** DNA_SIZE - 1) * (X_BOUND[1] - X_BOUND[0]) + X_BOUND[0]
    y = y_pop.dot(2 ** np.arange(DNA_SIZE)[::-1]) / float(2 ** DNA_SIZE - 1) * (Y_BOUND[1] - Y_BOUND[0]) + Y_BOUND[0]
    return x, y


def crossover_and_mutation(pop, CROSSOVER_RATE=0.8):
    new_pop = []
    for father in pop:  # 遍历种群中的每一个个体，将该个体作为父亲
        child = father  # 孩子先得到父亲的全部基因（这里我把一串二进制串的那些0，1称为基因）
        if np.random.rand() < CROSSOVER_RATE:  # 产生子代时不是必然发生交叉，而是以一定的概率发生交叉
            mother = pop[np.random.randint(POP_SIZE)]  # 再种群中选择另一个个体，并将该个体作为母亲
            cross_points = np.random.randint(low=0, high=DNA_SIZE * 2)  # 随机产生交叉的点
            child[cross_points:] = mother[cross_points:]  # 孩子得到位于交叉点后的母亲的基因
        mutation(child)  # 每个后代有一定的机率发生变异
        new_pop.append(child)

    return new_pop


def mutation(child, MUTATION_RATE=0.003):
    if np.random.rand() < MUTATION_RATE:  # 以MUTATION_RATE的概率进行变异
        mutate_point = np.random.randint(0, DNA_SIZE)  # 随机产生一个实数，代表要变异基因的位置
        child[mutate_point] = child[mutate_point] ^ 1  # 将变异点的二进制为反转


def select(pop, fitness):  # nature selection wrt pop's fitness
    idx = np.random.choice(np.arange(POP_SIZE), size=POP_SIZE, replace=True,
                           p=(fitness) / (fitness.sum()))
    return pop[idx]


def print_info(pop):
    fitness = get_fitness(pop)
    max_fitness_index = np.argmax(fitness)
    print("max_fitness:", fitness[max_fitness_index])
    x, y = translateDNA(pop)
    print("最优的基因型：", pop[max_fitness_index])
    print("(x, y):", (x[max_fitness_index], y[max_fitness_index]))
    print(F(x[max_fitness_index], y[max_fitness_index]))


if __name__ == "__main__":
    # 假设真值和估计值
    tru = np.array([1.348, 270.026, 93.867, 5.586, -0.918, 13.916])
    first_pre = np.array([1.38900e+00, 2.70347e+02, 9.39120e+01, 5.61600e+00, 2.63000e-01, 1.37100e+01])
    # 以估计值为中心，按±1.5范围生成初始种群,假设population_size=50
    population_size = 50
    error_range = np.array([0.177407, 0.33998787, 0.24179672, 0.34689838, 1.4670988, 0.44524276])
    all_population = []
    for i in range(6):
        single_dim = np.random.uniform(first_pre[i] - error_range[i], first_pre[i] + error_range[i], population_size)
        all_population.append(single_dim)
    all_population = np.array(all_population).T
    for _ in range(N_GENERATIONS):  # 迭代N代
        x, y = translateDNA(all_population)
        if 'sca' in locals():
            sca.remove()
        sca = ax.scatter(x, y, F(x, y), c='black', marker='o')
        plt.show()
        plt.pause(0.1)
        pop = np.array(crossover_and_mutation(pop, CROSSOVER_RATE))
        fitness = get_fitness(pop)
        pop = select(pop, fitness)  # 选择生成新的种群

    print_info(pop)
    plt.ioff()
    plot_3d(ax)
